// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/expression/expr_to_subfield_filter.h>
#include <gtest/gtest.h>
#include <pollux/testing/gtest_utils.h>
#include <pollux/expression/expr.h>
#include <pollux/functions/lib/is_null.h>
#include <pollux/functions/prestosql/registration/registration_functions.h>
#include <pollux/parse/expressions.h>
#include <pollux/parse/expressions_parser.h>
#include <pollux/parse/type_resolver.h>

namespace kumo::pollux::exec {
namespace {

using namespace kumo::pollux::common;

void validateSubfield(
    const Subfield& subfield,
    const std::vector<std::string>& expectedPath) {
  ASSERT_EQ(subfield.path().size(), expectedPath.size());
  for (int i = 0; i < expectedPath.size(); ++i) {
    ASSERT_TRUE(subfield.path()[i]);
    ASSERT_EQ(*subfield.path()[i], Subfield::NestedField(expectedPath[i]));
  }
}

class ExprToSubfieldFilterTest : public testing::Test {
 public:
  static void SetUpTestSuite() {
    functions::prestosql::registerAllScalarFunctions();
    parse::registerTypeResolver();
    memory::MemoryManager::testingSetInstance({});
  }

  core::TypedExprPtr parseExpr(
      const std::string& expr,
      const RowTypePtr& type) {
    return core::Expressions::inferTypes(
        parse::parseExpr(expr, {}), type, pool_.get());
  }

  core::CallTypedExprPtr parseCallExpr(
      const std::string& expr,
      const RowTypePtr& type) {
    auto call = std::dynamic_pointer_cast<const core::CallTypedExpr>(
        parseExpr(expr, type));
    POLLUX_CHECK_NOT_NULL(call);
    return call;
  }

  core::ExpressionEvaluator* evaluator() {
    return &evaluator_;
  }

 private:
  std::shared_ptr<memory::MemoryPool> pool_ =
      memory::memoryManager()->addLeafPool();
  std::shared_ptr<core::QueryCtx> queryCtx_{core::QueryCtx::create()};
  SimpleExpressionEvaluator evaluator_{queryCtx_.get(), pool_.get()};
};

TEST_F(ExprToSubfieldFilterTest, eq) {
  auto call = parseCallExpr("a = 42", ROW({{"a", BIGINT()}}));
  Subfield subfield;
  auto filter =
      ExprToSubfieldFilterParser::getInstance()->leafCallToSubfieldFilter(
          *call, subfield, evaluator());
  ASSERT_TRUE(filter);
  validateSubfield(subfield, {"a"});
  auto bigintRange = dynamic_cast<BigintRange*>(filter.get());
  ASSERT_TRUE(bigintRange);
  ASSERT_EQ(bigintRange->lower(), 42);
  ASSERT_EQ(bigintRange->upper(), 42);
  ASSERT_FALSE(bigintRange->testNull());
}

TEST_F(ExprToSubfieldFilterTest, eqExpr) {
  auto call = parseCallExpr("a = 21 * 2", ROW({{"a", BIGINT()}}));
  Subfield subfield;
  auto filter =
      ExprToSubfieldFilterParser::getInstance()->leafCallToSubfieldFilter(
          *call, subfield, evaluator());
  ASSERT_TRUE(filter);
  validateSubfield(subfield, {"a"});
  auto bigintRange = dynamic_cast<BigintRange*>(filter.get());
  ASSERT_TRUE(bigintRange);
  ASSERT_EQ(bigintRange->lower(), 42);
  ASSERT_EQ(bigintRange->upper(), 42);
  ASSERT_FALSE(bigintRange->testNull());
}

TEST_F(ExprToSubfieldFilterTest, eqSubfield) {
  auto call = parseCallExpr("a.b = 42", ROW({{"a", ROW({{"b", BIGINT()}})}}));
  Subfield subfield;
  auto filter =
      ExprToSubfieldFilterParser::getInstance()->leafCallToSubfieldFilter(
          *call, subfield, evaluator());
  ASSERT_TRUE(filter);
  validateSubfield(subfield, {"a", "b"});
  auto bigintRange = dynamic_cast<BigintRange*>(filter.get());
  ASSERT_TRUE(bigintRange);
  ASSERT_EQ(bigintRange->lower(), 42);
  ASSERT_EQ(bigintRange->upper(), 42);
  ASSERT_FALSE(bigintRange->testNull());
}

TEST_F(ExprToSubfieldFilterTest, neq) {
  auto call = parseCallExpr("a <> 42", ROW({{"a", BIGINT()}}));
  Subfield subfield;
  auto filter =
      ExprToSubfieldFilterParser::getInstance()->leafCallToSubfieldFilter(
          *call, subfield, evaluator());
  ASSERT_TRUE(filter);
  validateSubfield(subfield, {"a"});
  ASSERT_TRUE(filter->testInt64(41));
  ASSERT_FALSE(filter->testInt64(42));
  ASSERT_TRUE(filter->testInt64(43));
}

TEST_F(ExprToSubfieldFilterTest, lte) {
  auto call = parseCallExpr("a <= 42", ROW({{"a", BIGINT()}}));
  Subfield subfield;
  auto filter =
      ExprToSubfieldFilterParser::getInstance()->leafCallToSubfieldFilter(
          *call, subfield, evaluator());
  ASSERT_TRUE(filter);
  validateSubfield(subfield, {"a"});
  ASSERT_TRUE(filter->testInt64(41));
  ASSERT_TRUE(filter->testInt64(42));
  ASSERT_FALSE(filter->testInt64(43));
}

TEST_F(ExprToSubfieldFilterTest, lt) {
  auto call = parseCallExpr("a < 42", ROW({{"a", BIGINT()}}));
  Subfield subfield;
  auto filter =
      ExprToSubfieldFilterParser::getInstance()->leafCallToSubfieldFilter(
          *call, subfield, evaluator());
  ASSERT_TRUE(filter);
  validateSubfield(subfield, {"a"});
  ASSERT_TRUE(filter->testInt64(41));
  ASSERT_FALSE(filter->testInt64(42));
  ASSERT_FALSE(filter->testInt64(43));
}

TEST_F(ExprToSubfieldFilterTest, gte) {
  auto call = parseCallExpr("a >= 42", ROW({{"a", BIGINT()}}));
  Subfield subfield;
  auto filter =
      ExprToSubfieldFilterParser::getInstance()->leafCallToSubfieldFilter(
          *call, subfield, evaluator());
  ASSERT_TRUE(filter);
  validateSubfield(subfield, {"a"});
  ASSERT_FALSE(filter->testInt64(41));
  ASSERT_TRUE(filter->testInt64(42));
  ASSERT_TRUE(filter->testInt64(43));
}

TEST_F(ExprToSubfieldFilterTest, gt) {
  auto call = parseCallExpr("a > 42", ROW({{"a", BIGINT()}}));
  Subfield subfield;
  auto filter =
      ExprToSubfieldFilterParser::getInstance()->leafCallToSubfieldFilter(
          *call, subfield, evaluator());
  ASSERT_TRUE(filter);
  validateSubfield(subfield, {"a"});
  ASSERT_FALSE(filter->testInt64(41));
  ASSERT_FALSE(filter->testInt64(42));
  ASSERT_TRUE(filter->testInt64(43));
}

TEST_F(ExprToSubfieldFilterTest, between) {
  auto call = parseCallExpr("a between 40 and 42", ROW({{"a", BIGINT()}}));
  Subfield subfield;
  auto filter =
      ExprToSubfieldFilterParser::getInstance()->leafCallToSubfieldFilter(
          *call, subfield, evaluator());
  ASSERT_TRUE(filter);
  validateSubfield(subfield, {"a"});
  for (int i = 39; i <= 43; ++i) {
    ASSERT_EQ(filter->testInt64(i), 40 <= i && i <= 42);
  }
}

TEST_F(ExprToSubfieldFilterTest, in) {
  auto call = parseCallExpr("a in (40, 42)", ROW({{"a", BIGINT()}}));
  Subfield subfield;
  auto filter =
      ExprToSubfieldFilterParser::getInstance()->leafCallToSubfieldFilter(
          *call, subfield, evaluator());
  ASSERT_TRUE(filter);
  validateSubfield(subfield, {"a"});
  for (int i = 39; i <= 43; ++i) {
    ASSERT_EQ(filter->testInt64(i), i == 40 || i == 42);
  }
}

TEST_F(ExprToSubfieldFilterTest, isNull) {
  auto call = parseCallExpr("a is null", ROW({{"a", BIGINT()}}));
  Subfield subfield;
  auto filter =
      ExprToSubfieldFilterParser::getInstance()->leafCallToSubfieldFilter(
          *call, subfield, evaluator());
  ASSERT_TRUE(filter);
  validateSubfield(subfield, {"a"});
  ASSERT_FALSE(filter->testInt64(0));
  ASSERT_FALSE(filter->testInt64(42));
  ASSERT_TRUE(filter->testNull());
}

TEST_F(ExprToSubfieldFilterTest, isNotNull) {
  auto call = parseCallExpr("a is not null", ROW({{"a", BIGINT()}}));
  auto [subfield, filter] = toSubfieldFilter(call, evaluator());
  ASSERT_TRUE(filter);
  validateSubfield(subfield, {"a"});
  ASSERT_TRUE(filter->testInt64(0));
  ASSERT_TRUE(filter->testInt64(42));
  ASSERT_FALSE(filter->testNull());
}

TEST_F(ExprToSubfieldFilterTest, like) {
  auto call = parseCallExpr("a like 'foo%'", ROW({{"a", VARCHAR()}}));
  Subfield subfield;
  auto filter =
      ExprToSubfieldFilterParser::getInstance()->leafCallToSubfieldFilter(
          *call, subfield, evaluator());
  ASSERT_FALSE(filter);
}

TEST_F(ExprToSubfieldFilterTest, nonConstant) {
  auto call =
      parseCallExpr("a = b + 1", ROW({{"a", BIGINT()}, {"b", BIGINT()}}));
  Subfield subfield;
  auto filter =
      ExprToSubfieldFilterParser::getInstance()->leafCallToSubfieldFilter(
          *call, subfield, evaluator());
  ASSERT_FALSE(filter);
}

TEST_F(ExprToSubfieldFilterTest, userError) {
  auto call = parseCallExpr("a = 1 / 0", ROW({{"a", BIGINT()}}));
  Subfield subfield;
  auto filter =
      ExprToSubfieldFilterParser::getInstance()->leafCallToSubfieldFilter(
          *call, subfield, evaluator());
  ASSERT_FALSE(filter);
}

TEST_F(ExprToSubfieldFilterTest, dereferenceWithEmptyField) {
  auto call = std::make_shared<core::CallTypedExpr>(
      BOOLEAN(),
      std::vector<core::TypedExprPtr>{
          std::make_shared<core::DereferenceTypedExpr>(
              REAL(),
              std::make_shared<core::FieldAccessTypedExpr>(
                  ROW({{"", DOUBLE()}, {"", REAL()}, {"", BIGINT()}}),
                  std::make_shared<core::InputTypedExpr>(ROW(
                      {{"c0",
                        ROW({{"", DOUBLE()}, {"", REAL()}, {"", BIGINT()}})}})),
                  "c0"),
              1)},
      "is_null");
  Subfield subfield;
  auto filter =
      ExprToSubfieldFilterParser::getInstance()->leafCallToSubfieldFilter(
          *call, subfield, evaluator());
  ASSERT_FALSE(filter);
}

class CustomExprToSubfieldFilterParser : public ExprToSubfieldFilterParser {
 public:
  std::unique_ptr<common::Filter> leafCallToSubfieldFilter(
      const core::CallTypedExpr& call,
      common::Subfield& subfield,
      core::ExpressionEvaluator* evaluator,
      bool negated) override {
    if (call.inputs().empty()) {
      return nullptr;
    }

    const auto* leftSide = call.inputs()[0].get();

    if (call.name() == "custom_eq") {
      if (toSubfield(leftSide, subfield)) {
        return negated ? makeNotEqualFilter(call.inputs()[1], evaluator)
                       : makeEqualFilter(call.inputs()[1], evaluator);
      }
    } else if (call.name() == "is_null") {
      if (toSubfield(call.inputs()[0].get(), subfield)) {
        if (negated) {
          return isNotNull();
        }
        return isNull();
      }
    }
    return nullptr;
  }
};

class CustomExprToSubfieldFilterTest : public ExprToSubfieldFilterTest {
 public:
  static void SetUpTestSuite() {
    functions::prestosql::registerAllScalarFunctions("custom_");
    functions::registerIsNullFunction("is_null");
    parse::registerTypeResolver();
    memory::MemoryManager::testingSetInstance({});
    ExprToSubfieldFilterParser::registerParserFactory(
        []() { return std::make_unique<CustomExprToSubfieldFilterParser>(); });
  }

  static void TearDownTestSuite() {
    ExprToSubfieldFilterParser::registerParserFactory(
        []() { return std::make_unique<PrestoExprToSubfieldFilterParser>(); });
  }
};

TEST_F(CustomExprToSubfieldFilterTest, isNull) {
  auto call = parseCallExpr("a is null", ROW({{"a", BIGINT()}}));
  auto [subfield, filter] = toSubfieldFilter(call, evaluator());
  ASSERT_TRUE(filter);
  validateSubfield(subfield, {"a"});
  ASSERT_FALSE(filter->testInt64(0));
  ASSERT_FALSE(filter->testInt64(42));
  ASSERT_TRUE(filter->testNull());
}

TEST_F(CustomExprToSubfieldFilterTest, eq) {
  auto call = parseCallExpr("custom_eq(a, 42)", ROW({{"a", BIGINT()}}));
  auto [subfield, filter] = toSubfieldFilter(call, evaluator());
  ASSERT_TRUE(filter);
  validateSubfield(subfield, {"a"});
  auto bigintRange = dynamic_cast<BigintRange*>(filter.get());
  ASSERT_TRUE(bigintRange);
  ASSERT_EQ(bigintRange->lower(), 42);
  ASSERT_EQ(bigintRange->upper(), 42);
  ASSERT_FALSE(bigintRange->testNull());
}

TEST_F(CustomExprToSubfieldFilterTest, unsupported) {
  auto call = parseCallExpr("custom_neq(a, 42)", ROW({{"a", BIGINT()}}));
  POLLUX_ASSERT_USER_THROW(
      toSubfieldFilter(call, evaluator()),
      "Unsupported expression for range filter");
}

} // namespace
} // namespace kumo::pollux::exec
